package com.ntd.study.algorithm.reduce.huffman;

import cn.hutool.core.codec.Base64;
import cn.hutool.core.lang.Console;
import com.ntd.study.algorithm.reduce.StudyReduceMes;

import java.nio.charset.Charset;
import java.util.*;

/**
 * 哈夫曼编码
 */
public class HuffmanCode {

    public static void main(String[] args) {
        // 字符串转为byte数组
        String content = StudyReduceMes.mes;
        byte[] bytes = content.getBytes(Charset.defaultCharset());
        // 哈夫曼压缩
        byte[] huffmanCodes = huffmanCompress(bytes);
        final String encode = Base64.encode(huffmanCodes);
        Console.log("压缩后的字符串", encode);
        Console.log("压缩后的字符串长度", encode.length());
        // 哈夫曼解压缩
        byte[] decodeResult = decompress(codeTable, Base64.decode(encode));
        System.out.println(new String(decodeResult));
    }

    /**
     * 编码表
     */
    private static Map<Byte, String> codeTable;

    /**
     * 暂存可能不足一个字节的字符编码
     */
    private static String lastByte = "";

    /**
     * 将原始byte数组进行哈夫曼压缩
     * @param bytes 原始数组
     * @return 经过哈夫曼压缩的byte数组
     */
    public static byte[] huffmanCompress(byte[] bytes) {
        // step1 根据序列构建哈夫曼树
        List<HuffmanTreeNode> nodeList = getNodes(bytes);
        HuffmanTreeNode root = HuffmanTree.createHuffTree(nodeList);
        // step2 获取哈夫曼编码表
        codeTable = HuffmanTree.getCodingTable(root);
        // step3 根据编码表压缩byte数组
        return transfer(bytes, codeTable);
    }


    /**
     * 将byte数组转为哈夫曼树结点的list
     * @param bytes byte数组
     * @return 结点list
     */
    private static List<HuffmanTreeNode> getNodes(byte[] bytes) {
        // 字符与出现次数（即哈夫曼树的权）的映射
        Map<Byte, Integer> map = new HashMap<>(32);
        // 遍历bytes，将字符与其出现次数放入map
        for (byte b : bytes) {
            map.merge(b, 1, Integer::sum);
        }

        // 把map中的键值对转换成node加入到list中
        List<HuffmanTreeNode> nodeList = new ArrayList<>();
        for (Map.Entry<Byte, Integer> entry : map.entrySet()) {
            nodeList.add(new HuffmanTreeNode(entry.getKey(), entry.getValue()));
        }
        return nodeList;
    }


    /**
     * 将原始byte数组根据编码表转换为哈夫曼编码
     * @param codeTable 编码表
     * @param origin    原始数组
     * @return 哈夫曼编码后的byte数组
     */
    private static byte[] transfer(byte[] origin, Map<Byte, String> codeTable) {
        System.out.println(Arrays.toString(origin));
        // 将原数组转换成Huffman编码字符
        StringBuilder huffmanStr = new StringBuilder();
        for (byte item : origin) {
            huffmanStr.append(codeTable.get(item));
        }
        // 压缩数组的长度
        int len = huffmanStr.length() / 8;
        if (huffmanStr.length() % 8 != 0) {
            // 编码长度不能被8整除，保存多余的不足8位的部分
            lastByte = huffmanStr.substring(len * 8);
        }
        System.out.printf("lastByte = %s\n", lastByte);
        // 压缩存储的byte数组
        byte[] huffmanCode = new byte[len];
        for (int i = 0; i < huffmanCode.length; i++) {
            huffmanCode[i] = (byte) Integer.parseInt(huffmanStr.substring(i * 8, i * 8 + 8), 2);
        }
        System.out.printf("压缩后的长度为 %d\n", "".equals(lastByte) ? len : len + 1);
        return huffmanCode;
    }


    /**
     * 封装解压缩的方法
     *
     * @param bytes 经过压缩的byte数组
     * @return 解压后的byte数组
     */
    public static byte[] decompress(byte[] bytes) {
        return decompress(codeTable, bytes);
    }


    /**
     * 将一个byte转为二进制字符串
     *
     * @param b 一个byte
     * @return b对应的二进制字符串（补码）
     */
    private static String byteToString(byte b) {
        // 使用int暂存b
        int temp = b;
        // 补高位，和256（1 0000 0000）进行按位或
        temp |= 256;
        String str = Integer.toBinaryString(temp);
        return str.substring(str.length() - 8);
    }


    /**
     * 对压缩数据进行解码
     *
     * @param codeTable   编码表
     * @param huffmanCode 压缩的数组
     * @return 原字符串对应的byte数组
     */
    private static byte[] decompress(Map<Byte, String> codeTable, byte[] huffmanCode) {
        System.out.println(Arrays.toString(huffmanCode));
        // 先得到压缩数组对应的编码字符串
        StringBuilder stringBuilder = new StringBuilder();
        for (byte value : huffmanCode) {
            stringBuilder.append(byteToString(value));
        }
        stringBuilder.append(lastByte);
        // 把字符串按编码表进行解码(先将编码表反转，根据编码找原数值)
        Map<String, Byte> map = new HashMap<>(32);
        for (Map.Entry<Byte, String> entry : codeTable.entrySet()) {
            map.put(entry.getValue(), entry.getKey());
        }
        // 扫描StringBuilder
        List<Byte> byteList = new ArrayList<>();
        /*
         * 此处i借助count这个增量来移动
         * 若在每一轮循环时i++会导致漏扫描字符
         */
        for (int i = 0; i < stringBuilder.length(); ) {
            int count = 1;
            Byte b;
            while (true) {
                String str = stringBuilder.substring(i, i + count);
                b = map.get(str);
                if (b != null) {
                    byteList.add(b);
                    i += count;
                    break;
                }
                count++;
            }
        }
        // 将list赋值给array
        byte[] bytes = new byte[byteList.size()];
        for (int i = 0; i < bytes.length; i++) {
            bytes[i] = byteList.get(i);
        }
        return bytes;
    }
}
